from langchain.agents import tool
import pymysql
import os
from dotenv import load_dotenv
load_dotenv()


class MysqlTool:

    def __init__(self):
        self.connection = self.s_connection_db()

    @staticmethod
    def s_connection_db(self):
        if self.connection is None:
            # 建立数据库连接
            self.connection = pymysql.connect(
                host=os.getenv("db_host"),
                user=os.getenv("db_user"),
                password=os.getenv("db_password"),
                database=os.getenv("db_name")
            )
        return self.connection

    @tool
    def query_employee(self, query: dict):
        """只有需要查询客户信息时才会使用这个工具。请根据用户的输入，获取到关键字如: 名字或者邮箱"""
        # 建立数据库连接
        # connection = self.s_connection_db()

        # 创建游标对象
        cursor = self.connection.cursor()

        id = 0
        try:
            id = int(query.id) or 0
        except ValueError:
            id = 0

        # 执行SQL查询语句
        sql = f"SELECT * FROM customer where id = {id} or name = '{query.name}' or emp_card_no = '{query.emp_card_no}'"
        print("sql: ", sql)
        cursor.execute(sql)

        # 获取查询结果
        results = cursor.fetchall()

        info = ""
        # 处理查询结果
        for row in results:
            gender = '男' if row[2] == 1 else '女'
            info += f"id: {row[0]}, 姓名: {row[1]}, 性别: {gender}, 身份证号: {row[3]}, 工号: {row[4]}, 邮箱: {row[5]} \n"

        # 关闭游标对象和数据库连接
        cursor.close()

        return info